import os
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_milvus import Milvus
from dotenv import load_dotenv, find_dotenv
from langchain_openai import ChatOpenAI
from langchain_elasticsearch import ElasticsearchRetriever
from langchain.retrievers import EnsembleRetriever
from typing import Any, Dict, Mapping
from langchain_core.documents import Document
from flask import Flask, request
from common.json_response import JsonResponse
import pprint
_ = load_dotenv(find_dotenv())
api_key = os.environ["API_key"]
token = os.environ["Token"]

server = Flask(__name__)


def bm25_query(search_query: str) -> Dict:
    return {
        "query": {
            "match": {
                "content": search_query,
            },
        },
    }


def num_characters_mapper(hit: Mapping[str, Any]) -> Document:
    zhihu_link = hit["_source"]["zhihuLink"]
    metadata_question = hit["_source"]["question"]
    metadata_title = hit["_source"]["title"]
    page_content = hit["_source"]["content"]

    return Document(
        page_content=page_content,
        metadata={"article_zhihu_link": zhihu_link, "article_question": metadata_question,
                  "article_title": metadata_title,
                  "article_content": page_content},
    )


def get_llm():
    return ChatOpenAI(
        temperature=0.95,
        model="glm-4-plus",
        openai_api_key=api_key,
        openai_api_base="https://open.bigmodel.cn/api/paas/v4/")


@server.route("/api/get-quote", methods=['GET'])
def get_summarization():
    question = request.args.get('question')
    embeddings = ZhipuAIEmbeddings(
        model="embedding-3",
        api_key=api_key
    )
    vector_db = Milvus(
        embedding_function=embeddings,
        collection_name="excerpt_collection",
        connection_args={
            "uri": "http://api.caritas.pro:19530",
        },
        text_field="quote")
    vector_retriever = vector_db.as_retriever(search_kwargs={"k": 5})
    bm25_retriever = ElasticsearchRetriever.from_es_params(
        index_name="article_index",
        body_func=bm25_query,
        document_mapper=num_characters_mapper,
        url="http://api.caritas.pro:9200"
    )
    ensemble_retriever = EnsembleRetriever(
        retrievers=[bm25_retriever, vector_retriever], weights=[0.2, 0.8]
    )
    vector_documents = vector_retriever.invoke(question)
    pprint.pprint(vector_documents)
    documents = ensemble_retriever.invoke(question)
    document_dicts = [
        {
            "zhihu_link": doc.metadata["article_zhihu_link"],
            "question": doc.metadata["article_question"],
            "title": doc.metadata["article_title"],
            "content": doc.page_content
        }
        for doc in documents
    ]
    document_jsons = [doc_dict for doc_dict in document_dicts]
    return JsonResponse.success(document_jsons)


if __name__ == '__main__':
    server.run(host='0.0.0.0', port=5000)
